{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Question Answering with Memory\n",
    "\n",
    "This model shows a form of question answering with memory. You will bind two features (color and shape) by circular convolution and store them in a memory population. Then you will provide a cue to the model at a later time to determine either one of the features by deconvolution. This model exhibits better cognitive ability since the answers to the questions are provided at a later time and not at the same time as the questions themselves. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Setup the environment\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "\n",
    "import nengo\n",
    "from nengo.spa import Vocabulary\n",
    "from nengo.spa import Memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the model\n",
    "This model has parameters as described in the book, with `memory` population having 1000 neurons over 20 dimensions. The `memory` population is capable of storing a vector over time and it uses an integrator network to do so as discussed in the book. \n",
    "\n",
    "The color input presented to the model is `RED` and then `BLUE` for 0.25 seconds each before being turned off. In the same way the shape input is `CIRCLE` and then `SQUARE` for 0.25 seconds each. Thus, when you run the model, it will start by binding `RED` and `CIRCLE` for 0.25 seconds and then binding `BLUE` and `SQUARE` for 0.25 seconds. The cue for deconvolving bound semantic pointers will be turned off for 0.5 seconds and then starts cycling through `CIRCLE`, `RED`, `SQUARE`, and `BLUE` within each second. The model will be able to determine the correct answer using the cue, even when the color and shape inputs have been turned off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dim = 20       # Number of dimensions \n",
    "N_input = 300  # Number of neurons in population\n",
    "N_conv = 70    # Number of neurons per dimension in bind/unbind populations\n",
    "N_mem = 50     # Number of neurons per dimension in memory population\n",
    "\n",
    "# Creating the vocabulary\n",
    "rng = np.random.RandomState(7)  \n",
    "vocab = Vocabulary(dimensions=dim, rng=rng, max_similarity=0.1)\n",
    "\n",
    "model = nengo.Network(label='Question Answering with Memory', seed=12)  \n",
    "with model:\n",
    "    # Ensembles\n",
    "    A = nengo.Ensemble(n_neurons=N_input, dimensions=dim, label='A')\n",
    "    B = nengo.Ensemble(n_neurons=N_input, dimensions=dim, label='B')\n",
    "    C = nengo.Ensemble(n_neurons=N_input, dimensions=dim, label='C')\n",
    "    D = nengo.Ensemble(n_neurons=N_input, dimensions=dim, label='D')\n",
    "    E = nengo.Ensemble(n_neurons=N_input, dimensions=dim, label='E')\n",
    "    \n",
    "    # Creating memory population and connecting ensemble D to it\n",
    "    tau = 0.4\n",
    "    memory = nengo.networks.EnsembleArray(n_neurons=N_mem, n_ensembles=dim, label='Memory')\n",
    "    nengo.Connection(memory.output, memory.input, synapse=tau)     \n",
    "    nengo.Connection(D, memory.input, synapse=0.01)\n",
    "    \n",
    "    # Creating the Bind network\n",
    "    bind = nengo.networks.CircularConvolution(n_neurons=N_conv, dimensions=dim)\n",
    "    nengo.Connection(A, bind.A)\n",
    "    nengo.Connection(B, bind.B)\n",
    "    nengo.Connection(bind.output, D) \n",
    "    \n",
    "    # Creating the Unbind network\n",
    "    unbind = nengo.networks.CircularConvolution(n_neurons=N_conv, dimensions=dim, invert_a=True)\n",
    "    nengo.Connection(C, unbind.A)\n",
    "    nengo.Connection(memory.output, unbind.B)\n",
    "    nengo.Connection(unbind.output, E)\n",
    "    \n",
    "    # Getting semantic pointer values\n",
    "    CIRCLE = vocab.parse('CIRCLE').v\n",
    "    BLUE = vocab.parse('BLUE').v\n",
    "    RED = vocab.parse('RED').v\n",
    "    SQUARE = vocab.parse('SQUARE').v\n",
    "    ZERO = [0] * dim\n",
    "    \n",
    "    # function for providing color input\n",
    "    def color_input(t):\n",
    "        if t < 0.25:\n",
    "            return RED\n",
    "        elif t < 0.5:\n",
    "            return BLUE\n",
    "        else:\n",
    "            return ZERO\n",
    "\n",
    "    # function for providing shape input\n",
    "    def shape_input(t):\n",
    "        if t < 0.25:\n",
    "            return CIRCLE\n",
    "        elif t < 0.5:\n",
    "            return SQUARE\n",
    "        else:\n",
    "            return ZERO\n",
    "\n",
    "    # function for providing the cue\n",
    "    def cue_input(t):\n",
    "        if t < 0.5:\n",
    "            return ZERO\n",
    "        sequence = [ZERO, CIRCLE, RED, ZERO, SQUARE, BLUE]\n",
    "        idx = int(((t - 0.5) // (1. / len(sequence))) % len(sequence))\n",
    "        return sequence[idx]\n",
    "\n",
    "    # Defining inputs\n",
    "    inputA = nengo.Node(output=color_input, size_out=dim)\n",
    "    inputB = nengo.Node(output=shape_input, size_out=dim)\n",
    "    inputC = nengo.Node(output=cue_input, size_out=dim)\n",
    "    \n",
    "    # Connecting input to ensembles\n",
    "    nengo.Connection(inputA, A)\n",
    "    nengo.Connection(inputB, B)\n",
    "    nengo.Connection(inputC, C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add Probes to Collect Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with model:\n",
    "    A_probe = nengo.Probe(A, synapse=0.1)\n",
    "    B_probe = nengo.Probe(B, synapse=0.1)\n",
    "    C_probe = nengo.Probe(C, synapse=0.1)\n",
    "    D_probe = nengo.Probe(D, synapse=0.1)\n",
    "    E_probe = nengo.Probe(E, synapse=0.5)\n",
    "    memory_probe = nengo.Probe(memory.output, synapse=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with nengo.Simulator(model) as sim:  # Create the simulator\n",
    "    sim.run(3.0)                     # Run it for 3 seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "plt.subplot(6, 1, 1)\n",
    "plt.plot(sim.trange(), nengo.spa.similarity(sim.data[A_probe], vocab))\n",
    "plt.legend(vocab.keys, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=9)\n",
    "plt.ylabel(\"color\");\n",
    "\n",
    "plt.subplot(6, 1, 2)\n",
    "plt.plot(sim.trange(), nengo.spa.similarity(sim.data[B_probe], vocab))\n",
    "plt.legend(vocab.keys, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=9)\n",
    "plt.ylabel(\"shape\");\n",
    "\n",
    "plt.subplot(6, 1, 3)\n",
    "for pointer in ['RED * CIRCLE', 'BLUE * SQUARE']:\n",
    "    plt.plot(sim.trange(), vocab.parse(pointer).dot(sim.data[D_probe].T), label=pointer)\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0,fontsize=10)\n",
    "plt.ylabel(\"bound\");\n",
    "\n",
    "plt.subplot(6, 1, 4)\n",
    "for pointer in ['RED * CIRCLE', 'BLUE * SQUARE']:\n",
    "    plt.plot(sim.trange(), vocab.parse(pointer).dot(sim.data[memory_probe].T), label=pointer)\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0,fontsize=10)\n",
    "plt.ylabel(\"memory\");\n",
    "\n",
    "plt.subplot(6, 1, 5)\n",
    "plt.plot(sim.trange(), nengo.spa.similarity(sim.data[C_probe], vocab))\n",
    "plt.legend(vocab.keys, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=9)\n",
    "plt.ylabel(\"cue\");\n",
    "\n",
    "plt.subplot(6, 1, 6)\n",
    "plt.plot(sim.trange(), nengo.spa.similarity(sim.data[E_probe], vocab))\n",
    "plt.legend(vocab.keys, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, fontsize=9)\n",
    "plt.ylabel(\"Output\")\n",
    "plt.xlabel(\"time [s]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The graphs of `shape`, `color`, `bound` and `memory` show that first `RED * CIRCLE` and then `BLUE * SQUARE` are convolved (bound) and loaded into the `memory` population, so after 0.5 seconds the `memory` represents the superposition `RED * CIRCLE + BLUE * SQUARE`.\n",
    "\n",
    "The last plot shows that the output is most similar to the semantic pointer bound to the current cue. For example, when the cue is `CIRCLE`, the output is most similar to `RED`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the model using the `spa` package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import nengo\n",
    "import nengo.spa as spa\n",
    "from nengo.spa import Vocabulary\n",
    "import numpy as np\n",
    "\n",
    "D = 32  # the dimensionality of the vectors\n",
    "rng = np.random.RandomState(7)\n",
    "vocab = Vocabulary(dimensions=D, rng=rng, max_similarity=0.1)\n",
    "\n",
    "# Adding semantic pointers to the vocabulary\n",
    "CIRCLE = vocab.parse('CIRCLE')\n",
    "BLUE = vocab.parse('BLUE')\n",
    "RED = vocab.parse('RED')\n",
    "SQUARE = vocab.parse('SQUARE')\n",
    "ZERO = vocab.add('ZERO', [0]*D)\n",
    "\n",
    "model = spa.SPA(label=\"Question Answering with Memory\", vocabs=[vocab])\n",
    "with model:\n",
    "    model.A = spa.State(D)\n",
    "    model.B = spa.State(D)\n",
    "    model.C = spa.State(D)\n",
    "    model.D = spa.State(D)\n",
    "    model.E = spa.State(D)\n",
    "    model.memory = spa.State(D, feedback=1)\n",
    "\n",
    "    actions = spa.Actions(\n",
    "        'D = A * B',\n",
    "        'memory = D',\n",
    "        'E = memory * ~C'\n",
    "    )\n",
    " \n",
    "    model.cortical = spa.Cortical(actions)\n",
    "    \n",
    "    # function for providing color input\n",
    "    def color_input(t):\n",
    "        if t < 0.25:\n",
    "            return 'RED'\n",
    "        elif t < 0.5:\n",
    "            return 'BLUE'\n",
    "        else:\n",
    "            return 'ZERO'\n",
    "\n",
    "    # function for providing shape input\n",
    "    def shape_input(t):\n",
    "        if t < 0.25:\n",
    "            return 'CIRCLE'\n",
    "        elif t < 0.5:\n",
    "            return 'SQUARE'\n",
    "        else:\n",
    "            return 'ZERO'\n",
    "\n",
    "    # function for providing the cue\n",
    "    def cue_input(t):\n",
    "        if t < 0.5:\n",
    "            return 'ZERO'\n",
    "        sequence = ['ZERO', 'CIRCLE', 'RED', 'ZERO', 'SQUARE', 'BLUE']\n",
    "        idx = int(((t - 0.5) // (1. / len(sequence))) % len(sequence))\n",
    "        return sequence[idx]\n",
    "\n",
    "    # Inputs\n",
    "    model.input = spa.Input(A=color_input, B=shape_input, C=cue_input)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the model in nengo_gui"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from nengo_gui.ipython import IPythonViz\n",
    "IPythonViz(model, \"ch5-question-memory.py.cfg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Press the play button in the visualizer to run the simulation. You should see the \"semantic pointer cloud\" graphs as shown in the figure below.\n",
    "\n",
    "Graphs A, B, C show the colour, shape and cue inputs respectively. Graph E shows that the output is most similar to the semantic pointer which was initially bound to the given cue (showin in C). For example, when SQUARE is provided as a cue, the output is most similar to BLUE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "Image(filename='ch5-question-memory.png')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
